import os

from langchain_community.chat_models import ChatTongyi
from langchain_core.output_parsers import StrOutputParser, JsonOutputParser
from langchain_core.prompts import ChatPromptTemplate

os.environ["DASHSCOPE_API_KEY"] = "sk-9d8f1914800e497f8717144e860f99bc"
llm = ChatTongyi()

chunks = []
""""  直接使用LLM """
# for chunk in llm.stream("帮我制定一个健身计划"):
#     chunks.append(chunk)
#     print(chunk.content, end="|", flush=True)


""" 使用chain  """
# prompt = ChatPromptTemplate.from_template("tell me a joke about {topic}")
# parser = StrOutputParser()
# chains = prompt | llm | parser
# for chunk in chains.stream({"topic": "parrot"}):
#     print(chunk, end="|", flush=True)
#

"""  使用JSON进行解析  """
parser = JsonOutputParser()
chain = llm | parser
for chunk in chain.stream("output a list of the countries france, spain and japan and their populations in JSON format. "
    'Use a dict with an outer key of "countries" which contains a list of countries. '
    "Each country should have the key `name` and `population`"):
    print(chunk, flush=True)
